#!/usr/bin/env python
import os
import sys
from pymongo import MongoClient
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 获取MongoDB连接URI
mongo_uri = os.getenv('MONGO_URI', 'mongodb://localhost:27017/medical_system')

def update_appointments():
    """更新现有预约记录，添加缺失的状态字段"""
    client = MongoClient(mongo_uri)
    db = client.medical_system
    
    # 查询所有没有status字段的预约
    missing_status = db.appointments.find({"status": {"$exists": False}})
    status_count = 0
    
    for appointment in missing_status:
        db.appointments.update_one(
            {"_id": appointment["_id"]},
            {"$set": {"status": "预约中"}}
        )
        status_count += 1
    
    # 查询所有没有paymentStatus字段的预约
    missing_payment = db.appointments.find({"paymentStatus": {"$exists": False}})
    payment_count = 0
    
    for appointment in missing_payment:
        db.appointments.update_one(
            {"_id": appointment["_id"]},
            {"$set": {"paymentStatus": "未支付"}}
        )
        payment_count += 1
    
    print(f"成功为 {status_count} 条预约添加状态字段")
    print(f"成功为 {payment_count} 条预约添加支付状态字段")
    
    # 更新查询索引，确保可以按状态快速查询
    if "status_1" not in db.appointments.index_information():
        db.appointments.create_index("status")
        print("已创建status索引")
    
    if "paymentStatus_1" not in db.appointments.index_information():
        db.appointments.create_index("paymentStatus")
        print("已创建paymentStatus索引")

if __name__ == "__main__":
    print("开始更新预约记录...")
    update_appointments()
    print("更新完成!") 